package cc.shacocloud.luanniao.core.service.ai

import cc.shacocloud.luanniao.core.model.constant.AiModelTypeEnum
import cc.shacocloud.luanniao.core.model.constant.DocumentLanguageEnum
import cc.shacocloud.luanniao.core.service.ai.doubao.DoubaoAiModel
import cc.shacocloud.luanniao.core.service.ai.doubao.DoubaoAiOptions
import cc.shacocloud.luanniao.core.utils.messageFormat
import com.github.benmanes.caffeine.cache.Cache
import com.github.benmanes.caffeine.cache.Caffeine
import org.springframework.ai.chat.messages.SystemMessage
import org.springframework.ai.chat.messages.UserMessage
import java.util.concurrent.TimeUnit

/**
 * 生成 qa 格式响应的 提示词
 */

//     '<Task> The user will send a long text. Generate a Question and Answer pairs only using the knowledge in the long text. Please think step by step.'
//    'Step 1: Understand and summarize the main content of this text.\n'
//    'Step 2: What key information or concepts are mentioned in this text?\n'
//    'Step 3: Decompose or combine multiple pieces of information and concepts.\n'
//    'Step 4: Generate questions and answers based on these key information and concepts.\n'
//    '<Constraints> The questions should be clear and detailed, and the answers should be detailed and complete. '
//    'You must answer in {language}, in a style that is clear and detailed in {language}. No language other than {language} should be used. \n'
//    '<Format> Use the following format: Q1:\nA1:\nQ2:\nA2:...\n'
//    '<QA Pairs>'
const val GENERATOR_QA_PROMPT = """
    <任务> 用户将发送一条长文本。仅使用长文本中的知识生成问题和答案对。请一步步思考。
    步骤 1：理解并概括本文的主要内容。
    步骤 2：本文中提到了哪些关键信息或概念？
    步骤 3：分解或组合多条信息和概念。
    步骤 4：根据这些关键信息和概念生成问题和答案。
    <约束>问题应清晰、详细，答案应详细、完整。 
    您必须用{language}回答，并以{language}的方式清晰详细地回答。不得使用{language}以外的任何语言。 
    <格式> 使用以下格式：Q1:A1:Q2:A2:...
    <QA 对>
"""

val QA_REGEX =
    "Q\\d+[:：]\\s*(.*?)\\s*A\\d+[:：]\\s*([\\s\\S]*?)(?=Q\\d+[:：]|\$)".toRegex(
        setOf(
            RegexOption.IGNORE_CASE
        )
    )

val LINE_BREAKS_BLANKS_REGEX = "\\n\\s*".toRegex()

/**
 * 使用 ai 大模型进行 qa 文档的转换
 * @author 思追(shaco)
 */
fun AiModel.qaDocuments(
    text: String,
    language: DocumentLanguageEnum,
): List<QAPair> {

    val systemMessage = GENERATOR_QA_PROMPT.messageFormat(
        params = mapOf(
            "language" to language.language
        )
    )

    val result = chatCall(
        listOf(
            SystemMessage(systemMessage),
            UserMessage(text)
        )
    )

    // 拆分封装格式
    return QA_REGEX.findAll(result)
        .mapNotNull {
            val question = it.groups[1]?.value
            var answer = it.groups[2]?.value?.trim()
            answer = if (answer.isNullOrBlank()) null else LINE_BREAKS_BLANKS_REGEX.replace(answer, "\n")

            if (question.isNullOrBlank() || answer.isNullOrBlank()) {
                null
            } else {
                QAPair(question, answer)
            }
        }
        .toList()
}

data class QAPair(
    val question: String,
    val answer: String,
)

// ai 模型缓存
private val aiModelCache: Cache<String, AiModel> = Caffeine.newBuilder()
    .maximumSize(100)
    .expireAfterAccess(10, TimeUnit.MINUTES)
    .build()

/**
 * ai 模型属性 转为 [AiModel]
 */
fun AiOptions.toAiModel(): AiModel {
    return aiModelCache.get(getCacheKey()) {

        when (type) {
            AiModelTypeEnum.DOUBAO -> {
                this as DoubaoAiOptions
                DoubaoAiModel(apiKey, baseUrl, endpointId, defaultOptions)
            }
        }

    }
}